import sys
import os
import subprocess
import platform
import optparse
import cStringIO
import difflib

def main():
    # Run the test, capturing the output.
    test_module, options = parse_arguments()
    output = run_test_and_get_output(test_module)

    # Load the contents of the correct file.
    output_file_name = test_module + "_output.txt"
    correct_file_name = get_correct_file_name(test_module)
    with open(correct_file_name) as correct_file:
        correct_output = correct_file.read()

    # Compare the generated output with the correct output.
    if not options.graphical_diff:
        print_diffs(output, correct_output, output_file_name, correct_file_name)
    else:
        run_graphical_diff(
            output, correct_output, output_file_name, correct_file_name)

def parse_arguments():
    parser = optparse.OptionParser(usage="usage: %prog [options] test_module")
    parser.add_option("-g", "--graphical",
        action="store_true", dest="graphical_diff",
        default=False, help="Display a graphical diff of the results")
    options, args = parser.parse_args()

    if len(args) != 1:
        parser.error("You must specify a test module")

    return args[0], options

def run_test_and_get_output(test_module):
    # Redirect stdout and stderr.
    old_stdout, old_stderr = sys.stdout, sys.stderr
    new_stdout = cStringIO.StringIO()
    sys.stdout = new_stdout
    sys.stderr = sys.stdout

    # Run the test.
    try:
        __import__(test_module).main()
    except:
        sys.stdout, sys.stderr = old_stdout, old_stderr
        print new_stdout.getvalue()
        raise

    # Restore stdout and stderr and return the captured output.
    sys.stdout, sys.stderr = old_stdout, old_stderr
    return new_stdout.getvalue()

def get_correct_file_name(test_module):
    correct_file_name = test_module + "_correct.txt"

    # Look a platform-specific correct file.
    platform_specific_correct_file_name = (
        test_module + "_correct_" + platform.system().lower() + ".txt")
    if os.path.exists(platform_specific_correct_file_name):
        correct_file_name = platform_specific_correct_file_name

    return correct_file_name

def print_diffs(output, correct_output, output_file_name, correct_file_name):
    """Display a unified diff to the console."""
    output_lines = [
        line + "\n" for line in output.split("\n")]
    correct_output_lines = [
        line + "\n" for line in correct_output.split("\n")]

    for line in difflib.unified_diff(
            correct_output_lines, output_lines,
            correct_file_name, output_file_name):
        sys.stdout.write(line)

def run_graphical_diff(
        output, correct_output, output_file_name, correct_file_name):
    # Don't pop up the graphical diff if the output is the same.
    if output == correct_output:
        return

    # Save the output to disk so we can open it with diffuse.
    with open(output_file_name, "wb") as output_file:
        output_file.write(output)

    subprocess.Popen(
        "diffuse %s %s" % (output_file_name, correct_file_name), shell=True
        ).communicate()

    # Delete the temporary output file.
    os.unlink(output_file_name)

if __name__ == "__main__":
    main()

